// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/benchmark.h>
#include <melon/init/init.h>

#include <pollux/common/memory/memory.h>
#include <pollux/exec/set_accumulator.h>
#include <pollux/vector/vector_fuzzer.h>
#include <pollux/testing/vector/vector_test_base.h>

using namespace kumo;
using namespace kumo::pollux;
using namespace kumo::pollux::exec;

namespace {

// Adds 10M mostly unique values to a single SetAccumulator, then extracts
// unique values from it.
class SetAccumulatorBenchmark : public pollux::VectorBuilder {
 public:
  void setup() {
    VectorFuzzer::Options opts;
    opts.vectorSize = 1'000'000;
    VectorFuzzer fuzzer(opts, pool());

    auto rowType = ROW({"a", "b", "c"}, {BIGINT(), BIGINT(), VARCHAR()});
    for (auto i = 0; i < 10; ++i) {
      rowVectors_.emplace_back(fuzzer.fuzzInputRow(rowType));
    }
  }

  void runBigint() {
    runPrimitive<int64_t>("a");
  }

  void runVarchar() {
    runPrimitive<StringView>("c");
  }

  void runTwoBigints() {
    HashStringAllocator allocator(pool());
    const TypePtr type = ROW({BIGINT(), BIGINT()});
    aggregate::prestosql::SetAccumulator<ComplexType> accumulator(
        type, &allocator);

    for (const auto& row_vector : rowVectors_) {
      auto vector =
          make_row_vector({row_vector->childAt("a"), row_vector->childAt("b")});
      DecodedVector decoded(*vector);
      for (auto i = 0; i < row_vector->size(); ++i) {
        accumulator.addValue(decoded, i, &allocator);
      }
    }

    auto result = BaseVector::create(type, accumulator.size(), pool());
    accumulator.extractValues(*result, 0);
    melon::doNotOptimizeAway(result);
  }

 private:
  template <typename T>
  void runPrimitive(const std::string& name) {
    const auto& type = rowVectors_[0]->childAt(name)->type();

    HashStringAllocator allocator(pool());
    aggregate::prestosql::SetAccumulator<T> accumulator(type, &allocator);

    for (const auto& row_vector : rowVectors_) {
      DecodedVector decoded(*row_vector->childAt(name));
      for (auto i = 0; i < row_vector->size(); ++i) {
        accumulator.addValue(decoded, i, &allocator);
      }
    }

    auto result =
        BaseVector::create<FlatVector<T>>(type, accumulator.size(), pool());
    accumulator.extractValues(*result, 0);
    melon::doNotOptimizeAway(result);
  }

  std::vector<RowVectorPtr> rowVectors_;
};

std::unique_ptr<SetAccumulatorBenchmark> bm;

BENCHMARK(bigint) {
  bm->runBigint();
}

BENCHMARK(varchar) {
  bm->runVarchar();
}

BENCHMARK(twoBigints) {
  bm->runTwoBigints();
}

} // namespace

int main(int argc, char** argv) {
  melon::init(&argc, &argv);
  memory::MemoryManager::initialize({});

  bm = std::make_unique<SetAccumulatorBenchmark>();
  bm->setup();

  melon::runBenchmarks();

  bm.reset();

  return 0;
}
